import asyncio  # 用于异步编程，处理并发操作
import json  # 用于处理JSON格式的数据，存储和加载对话历史
import os  # 用于文件和目录操作
import sys  # 添加系统路径处理

# 将项目根目录添加到Python路径，确保可以导入项目中的其他模块
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# 从配置管理器导入各种配置获取函数
from config.config_manager import get_api_key, get_available_models, get_model, get_default_model
# 导入LLM API调用相关函数
from core.LLM_api import call_LLM_api, call_LLM_Tool_api
# 导入日志记录工具
from core.logger import get_logger

# 创建模块特定的日志记录器，方便跟踪调试
logger = get_logger("Conversation_Manager")

# 对话历史默认保存文件夹
CONVERSATIONS_DIR = "data/conversations"

# 定义终端颜色常量，用于美化命令行输出
CYAN = "\033[96m"    # 用户标签颜色
GREEN = "\033[92m"   # 助手标签颜色
RESET = "\033[0m"    # 重置颜色

class ConversationManager:
    """
    管理多轮对话的类，负责：
    1. 维护对话历史
    2. 处理用户追问
    3. 与LLM API交互获取响应
    4. 保存和加载对话历史
    """
    def __init__(self, system_message=None, model_name=None, conversations_dir=CONVERSATIONS_DIR, use_tools=False, mcp_config=None):
        """
        初始化对话管理器
        
        Args:
            system_message: 系统提示消息，设定AI助手的行为和角色
            model_name: 使用的模型名称，若为None则使用默认模型
            conversations_dir: 对话历史保存文件夹
            use_tools: 是否使用工具功能，如启用将允许模型调用外部工具
            mcp_config: MCP工具配置，包含工具服务器信息
        """
        self.messages = []  # 初始化对话历史列表
        
        # 从配置获取模型信息
        if model_name:
            selected_model = get_model(model_name=model_name)
            if not selected_model:
                logger.warning(f"未找到名称为 {model_name} 的模型，将使用默认模型")
                selected_model = get_model()  # 获取默认模型
        else:
            selected_model = get_model()  # 获取默认模型
        
        # 设置模型参数，如果配置获取失败则使用默认值
        if not selected_model:
            logger.error("无法获取有效的模型配置")
            self.model = "qwen3-235b-a22b"  # 默认备选模型
            self.api_key = ""
            self.base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
        else:
            self.model = selected_model.get("name")
            self.provider = selected_model.get("provider")
            self.base_url = selected_model.get("base_url")
            self.api_key = get_api_key(self.provider)  # 根据提供商获取API密钥
            
        logger.info(f"已配置模型: {self.model}, 提供商: {self.provider}")
            
        self.conversations_dir = conversations_dir
        self.use_tools = use_tools  # 是否启用工具功能标志
        self.mcp_config = mcp_config  # 工具配置信息
        
        # 确保对话历史文件夹存在
        self._ensure_conversations_dir()
        
        # 如果提供了系统消息，添加到对话历史
        if system_message:
            self.add_system_message(system_message)
            
        logger.info("对话管理器初始化完成")
    
    def set_model(self, model_name):
        """
        设置使用的模型，允许在运行时切换模型
        
        Args:
            model_name: 模型名称
        
        Returns:
            bool: 设置是否成功
        """
        selected_model = get_model(model_name=model_name)
        if not selected_model:
            logger.warning(f"未找到名称为 {model_name} 的模型，模型设置失败")
            return False
            
        self.model = selected_model.get("name")
        self.provider = selected_model.get("provider")
        self.base_url = selected_model.get("base_url")
        self.api_key = get_api_key(self.provider)
        logger.info(f"已切换模型: {self.model}, 提供商: {self.provider}")
        return True
    
    def get_available_models_info(self):
        """
        获取可用模型信息列表，用于UI显示和模型选择
        
        Returns:
            list: 包含模型信息的字典列表
        """
        return get_available_models()
        
    def add_system_message(self, content):
        """
        添加系统消息到对话历史，用于设置AI助手的行为和角色
        
        Args:
            content: 系统消息内容
        """
        self.messages.append({
            "role": "system",
            "content": content
        })
        logger.debug(f"添加系统消息: {content[:50]}..." if content else "无内容")
        
    def add_user_message(self, content):
        """
        添加用户消息到对话历史
        
        Args:
            content: 用户消息内容
        """
        self.messages.append({
            "role": "user",
            "content": content
        })
        logger.debug(f"添加用户消息: {content[:50]}..." if content else "无内容")
        
    def add_assistant_message(self, content):
        """
        添加助手消息到对话历史
        
        Args:
            content: 助手消息内容
        """
        self.messages.append({
            "role": "assistant",
            "content": content
        })
        logger.debug(f"添加助手消息: {content[:50]}..." if content else "无内容")
        
    async def get_response(self):
        """
        获取模型响应，处理最新的用户消息
        支持两种模式：基础API调用和工具增强API调用
        
        Returns:
            dict: 包含模型响应的字典，包括成功状态、内容、token计数等
        """
        logger.info("开始请求模型响应")
        
        if not self.api_key:
            logger.error("未找到API密钥，请在配置中设置")
            return {"success": False, "error_message": "未找到API密钥"}
            
        try:
            # 根据是否使用工具选择不同的API调用方法
            if self.use_tools and self.mcp_config:
                # 使用工具增强的API调用，允许模型调用外部工具
                result = await call_LLM_Tool_api(
                    self.api_key, 
                    self.messages, 
                    self.model, 
                    stream=True,  # 启用流式响应
                    mcp_config=self.mcp_config,
                    base_url=self.base_url  
                )
                
                # 从结果中提取响应内容
                content = result.get("full_content", "")
                success = True if content else False
                
                # 添加助手回复到对话历史
                if success:
                    self.add_assistant_message(content)
                    logger.info("成功获取工具增强模型响应并更新对话历史")
                
                return {
                    "success": success,
                    "content": content,
                    "reasoning_content": result.get("reasoning_content", ""),  # 模型的推理过程
                    "tool_calls": result.get("tool_calls", []),  # 工具调用记录
                    "total_prompt_tokens": result.get("total_prompt_tokens", 0),
                    "total_completion_tokens": result.get("total_completion_tokens", 0),
                    "api_call_count": result.get("api_call_count", 1)  # API调用次数，工具模式可能多次调用
                }
            else:
                # 使用基础API调用，不包含工具调用功能
                result = await call_LLM_api(
                    self.api_key, 
                    self.messages, 
                    self.model, 
                    stream=True,  # 启用流式响应
                    base_url=self.base_url
                )
                
                # 从结果中提取响应内容
                content = result.get("full_content", "")
                success = True if content else False
                
                # 添加助手回复到对话历史
                if success:
                    self.add_assistant_message(content)
                    logger.info("成功获取基础模型响应并更新对话历史")
                
                return {
                    "success": success,
                    "content": content,
                    "reasoning_content": result.get("reasoning_content", ""),
                    "prompt_tokens": result.get("prompt_tokens", 0),
                    "completion_tokens": result.get("completion_tokens", 0)
                }
                
        except Exception as e:
            error_message = str(e)
            logger.error(f"获取响应时出错: {error_message}")
            return {"success": False, "error_message": error_message}
    
    def get_conversation_history(self):
        """
        获取当前对话历史，用于UI显示或其他处理
        
        Returns:
            list: 对话历史列表
        """
        return self.messages
    
    def clear_history(self):
        """
        清除对话历史，但保留系统消息
        用于开始新对话但保持相同的AI助手角色设定
        """
        # 保存系统消息
        system_messages = [msg for msg in self.messages if msg["role"] == "system"]
        # 清除所有消息
        self.messages = []
        # 恢复系统消息
        self.messages.extend(system_messages)
        logger.info("对话历史已清除，保留系统消息")
    
    def _ensure_conversations_dir(self):
        """
        确保对话历史保存文件夹存在，不存在则创建
        这是一个内部辅助方法，不应被外部直接调用
        
        Returns:
            bool: 文件夹是否存在或创建成功
        """
        try:
            if not os.path.exists(self.conversations_dir):
                logger.info(f"创建对话历史保存文件夹: {self.conversations_dir}")
                os.makedirs(self.conversations_dir)
            return True
        except Exception as e:
            logger.error(f"创建对话历史文件夹失败: {str(e)}")
            return False
    
    def get_available_conversations(self):
        """
        获取可用的对话历史文件列表
        用于UI展示可加载的历史对话
        
        Returns:
            list: 对话历史文件名列表
        """
        if not os.path.exists(self.conversations_dir):
            return []
            
        try:
            # 获取所有.json文件
            files = [f for f in os.listdir(self.conversations_dir) 
                    if f.endswith('.json') and os.path.isfile(os.path.join(self.conversations_dir, f))]
            files.sort()  # 按文件名排序
            return files
        except Exception as e:
            logger.error(f"获取对话历史文件列表失败: {str(e)}")
            return []
    
    def save_conversation(self, file_name):
        """
        保存对话历史到文件，如果文件夹不存在则创建
        
        Args:
            file_name: 保存文件名
            
        Returns:
            bool: 保存是否成功
        """
        # 确保文件夹存在
        if not self._ensure_conversations_dir():
            return False
        
        # 构建完整文件路径
        file_path = os.path.join(self.conversations_dir, file_name)
        
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.messages, f, ensure_ascii=False, indent=2)  # 格式化JSON保存
            logger.info(f"对话历史已保存到: {file_path}")
            return True
        except Exception as e:
            logger.error(f"保存对话历史失败: {str(e)}")
            return False
    
    def load_conversation(self, file_name):
        """
        从文件加载对话历史
        
        Args:
            file_name: 加载文件名
            
        Returns:
            bool: 加载是否成功
        """
        # 构建完整文件路径
        file_path = os.path.join(self.conversations_dir, file_name)
        
        if not os.path.exists(file_path):
            logger.error(f"对话历史文件不存在: {file_path}")
            return False
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                self.messages = json.load(f)  # 从JSON加载对话历史
            logger.info(f"从 {file_path} 加载对话历史，共 {len(self.messages)} 条消息")
            return True
        except Exception as e:
            logger.error(f"加载对话历史失败: {str(e)}")
            return False


# 程序入口点：作为独立脚本运行时执行
if __name__ == '__main__':
    # 创建异步主函数
    async def main():
        # MCP配置：定义工具服务器配置
        mcp_config = {
            "mcpServers": {
                "modelscope_cantian-ai_Bazi-MCP": {  # 模型空间MCP服务
                    "url": "https://mcp.api-inference.modelscope.net/8d92e83e08df46/sse",
                },
                "GYCDWcq578JMgRTd34B0_": {  # 自定义seatable MCP服务器
                    "name": "自定义seatable MCP 服务器",
                    "type": "stdio",  # 标准输入输出类型
                    "description": "",
                    "isActive": True,
                    "registryUrl": "",
                    "command": "uv",  # 使用uv运行Python脚本
                    "args": [
                        "--directory",
                        r"F:\流量工作室\python_code\WeChat_AI_Assistant_1.1",  # 工作目录
                        "run", "tools/SeaTable_server.py"  # 运行SeaTable工具服务器
                    ]
                }
            }
        }
        
        # 设置AI助手的系统提示
        system_prompt = "你是一名智能助手，善于调用工具帮我解决问题。"

        # 询问是否使用工具
        use_tools = input("是否启用工具支持? (y/n): ").lower() == 'y'
        
        # 获取可用模型列表并显示
        available_models = get_available_models()
        default_model = get_default_model()
        
        if available_models:
            print("\n可用模型:")
            for model in available_models:
                default_mark = " (默认)" if model.get("id") == default_model.get("id") else ""
                print(f"{model.get('id')}. {model.get('name')} - {model.get('description')}{default_mark}")
                
            # 处理模型选择
            model_choice = input("\n请选择要使用的模型编号 (直接回车使用默认): ")
            
            selected_model = None
            if not model_choice:  # 用户直接回车，使用默认模型
                selected_model = default_model
            else:
                try:
                    model_id = int(model_choice)
                    for model in available_models:
                        if model.get("id") == model_id:
                            selected_model = model
                            break
                    
                    if not selected_model:
                        print(f"未找到ID为{model_id}的模型，将使用默认模型")
                        selected_model = default_model
                except ValueError:
                    print("输入的不是有效编号，将使用默认模型")
                    selected_model = default_model
                    
            model_name = selected_model.get("name")
            print(f"已选择模型: {model_name}")
        else:
            print("配置文件中未找到可用模型，将使用默认模型 qwen3-235b-a22b")
            model_name = "qwen3-235b-a22b"
        
        # 初始化对话管理器
        conversation = ConversationManager(
            system_message=system_prompt,                   # 设置AI助手的角色和行为特征
            model_name=model_name,                          # 用户选择的模型
            use_tools=use_tools,                            # 控制是否启用工具功能
            mcp_config=mcp_config if use_tools else None    # 工具配置，仅在启用工具时提供
        )
        
        # 检查是否有可用的历史对话
        available_conversations = conversation.get_available_conversations()
        
        # 如果有历史对话，询问用户是否加载
        if available_conversations:
            print(f"发现 {len(available_conversations)} 个历史对话文件:")
            for i, file_name in enumerate(available_conversations):
                print(f"{i+1}. {file_name}")
                
            load = input("\n是否加载历史对话? (y/n): ")
            if load.lower() == 'y':
                while True:
                    try:
                        choice = input("请输入要加载的对话编号 (输入0取消): ")
                        if choice == '0':
                            break
                            
                        idx = int(choice) - 1
                        if 0 <= idx < len(available_conversations):
                            if conversation.load_conversation(available_conversations[idx]):
                                print(f"成功加载对话: {available_conversations[idx]}")
                                break
                        else:
                            print("无效的选择，请重新输入")
                    except ValueError:
                        print("请输入有效的数字")
        
        print("\n多轮对话示例 (输入 'exit' 退出)")
        
        # 多轮对话循环
        while True:
            # 获取用户输入
            user_input = input(f"\n{CYAN}用户:{RESET} ")
            
            # 检查退出条件
            if user_input.lower() == 'exit':
                break
                
            # 添加用户消息并获取响应
            conversation.add_user_message(user_input)       # 将用户输入添加到对话历史
            response = await conversation.get_response()    # 异步获取模型响应
            
            # 显示助手回复
            if response["success"]:
                print(f"\n{GREEN}助手:{RESET} {response['content']}")
                
                # 如果使用工具，显示更多调试信息
                if use_tools and 'api_call_count' in response:
                    print(f"\nAPI调用次数: {response['api_call_count']}")
                    print(f"总输入Token: {response['total_prompt_tokens']}")
                    print(f"总输出Token: {response['total_completion_tokens']}")
            else:
                print(f"\n{GREEN}助手:{RESET} 抱歉，出现了错误: {response.get('error_message')}")
        
        # 保存对话历史（可选）
        save = input("\n是否保存对话历史? (y/n): ")
        if save.lower() == 'y':
            file_name = input("输入保存文件名: ")
            if not file_name.endswith('.json'):
                file_name += '.json'
            conversation.save_conversation(file_name)
    
    # 运行异步主函数
    asyncio.run(main()) 